"""
    model: iic/nlp_mt5_zero-shot-augment_chinese-base
    link: https://modelscope.cn/models/iic/nlp_mt5_zero-shot-augment_chinese-base

该模型在mt5模型基础上使用了大量中文数据进行训练，并引入了零样本分类增强的技术，使模型输出稳定性大幅提升。
    支持任务包含：
        文本分类：给定一段文本和候选标签，模型可输出文本所属的标签。
        自然语言推理：给定两段文本，判断两者关系。
        阅读理解：给定问题和参考文本，输出问题的答案。
        问题生成：给定答案和参考文本，生成该答案对应的问题。
        摘要生成：给定一段文本，生成该文本的摘要。
        标题生成：给定一段文本，为其生成标题。
        评价对象抽取：给定一段文本，抽取该段文本的评价对象。
        翻译：给定一段文本，将其翻译成另一种语言。
"""
import torch
import platform
from transformers import T5Tokenizer, T5ForConditionalGeneration, T5Config


# 获取当前操作系统的名称
os_name = platform.system()
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

# 设置模型路径及数据集路径
if os_name == "Windows":
    model_dir = r'D:\test_transformers01\stance\model\nlp_mt5_zero-shot-augment_chinese-base'
    print("当前执行环境是 Windows...")
elif os_name == "Linux":
    model_dir = r'/root/autodl-fs/models/nlp_mt5_zero-shot-augment_chinese-base'
    print("当前执行环境是 Linux...")
else:
    raise ValueError("当前执行环境不是 Windows 也不是 Linux")


# 1、加载tokenizer及预训练模型
tokenizer = T5Tokenizer.from_pretrained(model_dir, legacy=False)
model = T5ForConditionalGeneration.from_pretrained(model_dir)
model.to(device)


def t5inference(text, task_prefix, max_length=None, min_length=1):
    """
      :param text: 要生成摘要的文本
      :param task_prefix: 执行的任务
      :param max_length: 摘要的最大长度
      :return:
      """
    # 准备前缀+文本
    t5_prepared_text = f'{task_prefix}: ' + text

    print("input text: \n", t5_prepared_text)

    # 分词
    tokenized_text = tokenizer.encode(t5_prepared_text, return_tensors="pt").to(device)

    # 进行文本摘要
    # prepare_inputs_for_generation
    summary_ids = model.generate(tokenized_text,
                                 num_beams=4, # 束搜索法
                                 no_repeat_ngram_size=2, # 确保不重复
                                 min_length=min_length, # 长度限制
                                 max_length=max_length,
                                 early_stopping=True # 提前停止
    )
    # 将id转换为输出 summary_ids.shape = [1, length]
    output = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return output

def t0():
    text = """“足球从未高于生死”，这是3年前欧洲杯赛场上丹麦球员埃里克森心脏骤停时，各路媒体报道该事件用的最多的表达。
          而在经历了那段惊心动魄但又充满人情味的艰难时刻后，32岁的埃里克森时隔1100天再次为国征战欧洲杯，而且奉献了进球。
          17日凌晨的欧洲杯小组赛，埃里克森进球的那一刻，感动和欣慰扑面而来。最终丹麦和斯洛文尼亚队1比1战平，各取1分。

          丹麦队对垒斯洛文尼亚，这场热度并不算高的小组赛首轮争夺因为一个人的出现得到了外界的关注，他就是埃里克森。
          曼联中场在在第17分钟的进球帮助祖国球队取得了领先，他也在经历上届欧洲杯的心脏骤停的遭遇之后，实现了“王者归来”。
          尽管这次破门遗憾没能帮助丹麦队最终获得胜利，但绰号“爱神”的埃里克森依然得到了全场乃至全世界球迷的祝福。

        """
    output = t5inference(text, task_prefix='文本摘要', min_length=20, max_length=64)
    print('模型结果:\n', output)

    # text = """
    #     候选标签：故事,房产,娱乐,文化,游戏,国际,股票,科技,军事,教育。
    #     文本内容：他们的故事平静而闪光，一代人奠定沉默的基石，让中国走向繁荣。
    #     """
    # output = t5inference(text, task_prefix='文本分类', max_length=64)
    # print('模型结果:\n', output)
    #
    # text = """
    #         如果日本沉没，中国会接收日本难民吗？
    #     """
    # output = t5inference(text, task_prefix='翻译成英文', max_length=512)
    # print('模型结果:\n', output)


if __name__ == '__main__':
    # print(model.config)
    # print(model)
    t0()
